import pandas as pd
import re
import jieba

# 停用词表，适当添加更多停用词
STOPWORDS = set(["的", "了", "在", "是", "和", "我", "有", "也", "不", "都", "很", "这个", "就", "都", "人", "不是"])


# 清理无意义字符，保留中文
def clean_text(text):
    # 保留中文字符，去除其他字符
    text = re.sub(r'[^\u4e00-\u9fa5]', '', text)
    # 你也可以增加一些规则来去除特定词汇
    return text


# 去除停用词
def remove_stopwords(text):
    return ' '.join([word for word in jieba.cut(text) if word not in STOPWORDS])


def preprocess_data():
    # 读取数据
    df = pd.read_csv("dataset/ratings.csv")

    # 清理文本数据
    df['cleaned_comment'] = df['comment'].apply(lambda x: clean_text(str(x)))

    # 去除停用词
    df['processed_comment'] = df['cleaned_comment'].apply(remove_stopwords)

    # 生成标签：评分大于等于4为正面（1），评分小于等于2为负面（0），评分为3则不处理
    df['label'] = df['rating'].apply(lambda x: 1 if x >= 4 else (0 if x <= 2 else None))

    # 删除无效行（中性评价）
    df = df.dropna(subset=['label'])
    df['label'] = df['label'].astype(int)

    # 保存处理后的数据
    df[['processed_comment', 'label']].to_csv("processed_data.csv", index=False, header=False)

    print("数据处理完成。")
    print(f"总数据量：{len(df)}")


if __name__ == "__main__":
    preprocess_data()
